import pandas as pd
import numpy as np
import akshare as ak
import matplotlib.pyplot as plt
import mplfinance as mpf

from mplfinance.original_flavor import candlestick2_ohlc
from matplotlib.ticker import FormatStrFormatter


def getStoneDataAll(symbol):
    pingan = ak.stock_zh_a_daily(symbol=symbol, adjust="qfq")
    data = pingan.reset_index()
    data = data.dropna(how='any').reset_index(drop=True)  # 去除空值且从零开始编号索引
    data = data.sort_values(by='date', ascending=True)
    data['date'] = pd.to_datetime(data['date'])
    data.set_index("date", inplace=True)
    data.drop(columns=['index'], inplace=True)
    # 均线数据
    data['1'] = data['close']
    data['5'] = data.close.rolling(5).mean()
    data['10'] = data.close.rolling(10).mean()
    # 生成买卖信号
    data['Signal'] = 0
    # 短期均线上穿长期均线，产生买入信号
    data.loc[data['5'] > data['10'], 'Signal'] = 1
    # 短期均线下穿长期均线，产生卖出信号
    data.loc[data['5'] < data['10'], 'Signal'] = -1
    return data

def getStoneData30(symbol):
    pingan = ak.stock_zh_a_daily(symbol=symbol, adjust="qfq")
    data = pingan.reset_index().iloc[-30:, :6]  # 取过去30天数据
    data = data.dropna(how='any').reset_index(drop=True)  # 去除空值且从零开始编号索引
    data = data.sort_values(by='date', ascending=True)
    data['date'] = pd.to_datetime(data['date'])
    data.set_index("date", inplace=True)
    data.drop(columns=['index'], inplace=True)
    # 均线数据
    data['1'] = data['close']
    data['5'] = data.close.rolling(5).mean()
    data['10'] = data.close.rolling(10).mean()
    # 生成买卖信号
    data['Signal'] = 0
    # 短期均线上穿长期均线，产生买入信号
    data.loc[data['5'] > data['10'], 'Signal'] = 1
    # 短期均线下穿长期均线，产生卖出信号
    data.loc[data['5'] < data['10'], 'Signal'] = -1
    return data


def getStoneData():
    pingan = ak.stock_zh_a_daily(symbol="sh601318", adjust="qfq")
    df3 = pingan.reset_index().iloc[-30:, :6]  # 取过去30天数据
    df3 = df3.dropna(how='any').reset_index(drop=True)  # 去除空值且从零开始编号索引
    df3 = df3.sort_values(by='date', ascending=True)
    # print(df3.info())

    # 均线数据
    df3['5'] = df3.close.rolling(5).mean()
    df3['10'] = df3.close.rolling(10).mean()

    # print(df3.tail())
    return df3


def drawStockTrend(df3):
    # plt.style.use("ggplot")
    # 显示中文
    # plt.rcParams['font.sans-serif'] = ['Microsoft YaHei']

    # print(df3.dtypes)
    df3['date'] = pd.to_datetime(df3['date'])
    symbol = "sh601318"
    df3.set_index("date", inplace=True)
    # print(df3.dtypes)
    # print(df3)
    # 绘制K线图
    # df3.drop(columns=['index'], inplace=True)
    # df3['1'] = df3['close'].shift(1)
    # mpf.plot(df3, type='candle', mav=(1, 5, 10), volume=True, mavcolors=['r', 'g', 'b'])
    # 计算1日线与5日线的交叉点
    # df3['yesterday'] = df3['close'].shift(1)
    # df3['crossover_up'] = (df3['open'] > df3['yesterday']) & (df3['open'].shift(4) <= df3['yesterday'])
    # df3['crossover_down'] = (df3['open'] < df3['yesterday']) & (df3['open'].shift(4) >= df3['yesterday'])
    # df3.drop(columns=['index'], inplace=True)
    # print(df3.columns)
    # print(df3[['close', 'open',  'yesterday', 'crossover_up',
    #            'crossover_down']])

    # # 绘制交叉点的标记点
    # plt.scatter(df3.index[df3['crossover_up']], df3['open'][df3['crossover_up']], color='r')
    # plt.scatter(df3.index[df3['crossover_down']], df3['open'][df3['crossover_down']], color='g')

    # 绘制交叉点的标记点，直接在现有的图表上添加
    plt.figure(figsize=(10, 6))
    # plt.scatter(df3.index[df3['crossover_up']], df3['open'][df3['crossover_up']] - df3['open'].shift(4), color='r', marker='^')
    # plt.scatter(df3.index[df3['crossover_down']], df3['open'][df3['crossover_down']] - df3['open'].shift(4), color='g', marker='v')

    plt.plot(df3['close'].values, alpha=0.5, label='MA1')
    plt.plot(df3['5'].values, alpha=0.5, label='MA5')
    plt.plot(df3['10'].values, alpha=0.5, label='MA10')

    # 画均线
    plt.plot(df3['5'].values, alpha=0.5, label='MA5')
    plt.plot(df3['10'].values, alpha=0.5, label='MA10')

    plt.title(f"{symbol} 中国平安股价走势")
    # 修改x轴坐标
    plt.xticks(ticks=np.arange(0, len(df3)), labels=df3.index.astype(str).to_numpy())
    plt.xticks(rotation=45, size=8)
    plt.xlabel("日期")
    plt.ylabel("价格")
    # x轴坐标显示不全，整理
    # plt.subplots_adjust(bottom=0.25)
    plt.show(block=True)


# 绘制日均线
def drawStoneTrend1(df3):
    plt.rcParams['font.sans-serif'] = ['Microsoft YaHei']
    df3['date'] = pd.to_datetime(df3['date'])
    symbol = "sh601318"
    df3.set_index("date", inplace=True)
    plt.figure(figsize=(10, 6))
    plt.plot(df3['1'].values, alpha=0.5, label='MA1')
    plt.plot(df3['5'].values, alpha=0.5, label='MA5')
    plt.plot(df3['10'].values, alpha=0.5, label='MA10')

    plt.title(f"{symbol} 中国平安股价走势")
    # 修改x轴坐标
    plt.xticks(ticks=np.arange(0, len(df3)), labels=df3.index.astype(str).to_numpy())
    plt.xticks(rotation=45, size=8)
    plt.xlabel("日期")
    plt.ylabel("价格")
    plt.legend()
    # x轴坐标显示不全，整理
    plt.subplots_adjust(bottom=0.25)
    plt.show(block=True)


def drawStroneTrend2(data):
    data['1'] = data['close']
    # 生成买卖信号
    data['Signal'] = 0
    print(data)

    # 短期均线上穿长期均线，产生买入信号
    data.loc[data['5'] > data['10'], 'Signal'] = 1
    # 短期均线下穿长期均线，产生卖出信号
    data.loc[data['5'] < data['10'], 'Signal'] = -1

    plt.rcParams['font.sans-serif'] = ['Microsoft YaHei']
    data['date'] = pd.to_datetime(data['date'])
    symbol = "sh601318"
    data.set_index("date", inplace=True)

    # 绘制股价和移动平均线
    plt.figure(figsize=(10, 6))
    plt.plot(data['1'], alpha=0.5, label='M1')
    plt.plot(data['5'], alpha=0.5, label='M5')
    plt.plot(data['10'], alpha=0.5, label='M10')

    # 标记买卖信号
    plt.scatter(data[data['Signal'] == 1].index, data[data['Signal'] == 1]['5'], marker='^', color='r',
                label='Buy Signal')
    plt.scatter(data[data['Signal'] == -1].index, data[data['Signal'] == -1]['10'], marker='v', color='g',
                label='Sell Signal')
    # plt.xticks(ticks=np.arange(0, len(data)), labels=data.index.astype(str).to_numpy())
    # plt.scatter(y=data[data['Signal'] == 1]['5'], marker='^', color='g', label='Buy Signal')
    # plt.scatter(y=data[data['Signal'] == -1]['10'], marker='v', color='r', label='Sell Signal')
    plt.xticks(rotation=45, size=8)
    plt.title(f"{symbol} 中国平安股价走势")
    plt.xlabel("日期")
    plt.ylabel("价格")
    plt.legend()
    # x轴坐标显示不全，整理
    plt.subplots_adjust(bottom=0.25)
    plt.show(block=True)


def drawStoneTrend3(symbol, data):
    plt.rcParams['font.sans-serif'] = ['Microsoft YaHei']
    # 绘制股价和移动平均线
    plt.figure(figsize=(10, 6))
    plt.plot(data['1'], alpha=0.5, label='M1')
    plt.plot(data['5'], alpha=0.5, label='M5')
    plt.plot(data['10'], alpha=0.5, label='M10')

    # 标记买卖信号
    plt.scatter(data[data['Signal'] == 1].index, data[data['Signal'] == 1]['5'], marker='^', color='r',
                label='Buy Signal')
    plt.scatter(data[data['Signal'] == -1].index, data[data['Signal'] == -1]['10'], marker='v', color='g',
                label='Sell Signal')
    # plt.xticks(ticks=np.arange(0, len(data)), labels=data.index.astype(str).to_numpy())
    # plt.scatter(y=data[data['Signal'] == 1]['5'], marker='^', color='g', label='Buy Signal')
    # plt.scatter(y=data[data['Signal'] == -1]['10'], marker='v', color='r', label='Sell Signal')
    plt.xticks(rotation=45, size=8)
    plt.title(f"{symbol} 中国平安股价走势")
    plt.xlabel("日期")
    plt.ylabel("价格")
    plt.legend()
    # x轴坐标显示不全，整理
    plt.subplots_adjust(bottom=0.25)
    plt.show(block=True)


# 绘制K线图
def drawKLine(df3):
    plt.style.use("ggplot")
    fig, ax = plt.subplots(1, 1, figsize=(8, 3), dpi=200)
    # 绘制 K线
    candlestick2_ohlc(ax,
                      opens=df3['open'].values,
                      highs=df3['high'].values,
                      lows=df3['low'].values,
                      closes=df3['close'].values,
                      width=0.75, colorup="r", colordown="g")

    # 显示最高点和最低点
    ax.text(df3.high.idxmax(), df3.high.max(), s=df3.high.max(), fontsize=8)
    ax.text(df3.high.idxmin(), df3.high.min() - 2, s=df3.high.min(), fontsize=8)
    # 显示中文
    plt.rcParams['font.sans-serif'] = ['Microsoft YaHei']

    ax.set_facecolor("white")
    ax.set_title("中国平安")

    # 画均线
    plt.plot(df3['5'].values, alpha=0.5, label='MA5')
    plt.plot(df3['10'].values, alpha=0.5, label='MA10')

    ax.legend(facecolor='white', edgecolor='white', fontsize=6)
    # date 为 object 数据类型，通过 pd.to_datetime将该列数据转换为时间类型，即datetime
    df3.date = pd.to_datetime(df3.date, format='%Y-%m-%d')
    # 修改x轴坐标
    plt.xticks(ticks=np.arange(0, len(df3)), labels=df3.date.dt.strftime('%Y-%m-%d').to_numpy())
    plt.xticks(rotation=90, size=8)
    # 修改y轴坐标
    ax.yaxis.set_major_formatter(FormatStrFormatter('%.2f'))
    # x轴坐标显示不全，整理
    plt.subplots_adjust(bottom=0.25)
    plt.show(block=True)


def backtest1(data):
    # 显示中文
    plt.rcParams['font.sans-serif'] = ['Microsoft YaHei']
    # 计算每日收益率
    data['Daily_Return'] = data['close'].pct_change()

    # 计算策略信号的收益率（shift(1) 是为了避免未来数据的偏差）
    data['Strategy_Return'] = data['Signal'].shift(1) * data['Daily_Return']

    # 计算累计收益
    data['Cumulative_Return'] = (1 + data['Strategy_Return']).cumprod()
    print(data)
    # 绘制累计收益曲线
    plt.figure(figsize=(10, 6))
    plt.plot(data['Cumulative_Return'], label='策略累计回报', color='b')
    plt.plot(data['close'] / data['close'].iloc[0], label='股票累计回报', color='g')
    plt.title("策略累计回报 vs. 股票累计回报")
    plt.xlabel("日期")
    plt.ylabel("策略累计回报")
    plt.legend()
    plt.show(block=True)


def backtest(data):
    # 初始化交叉信号列
    data['Signal'] = 0
    # 计算每日收益率
    data['Daily_Return'] = data['close'].pct_change()

    # 计算策略信号
    data['Signal'] = 0
    data.loc[data['Daily_Return'] > 0, 'Signal'] = 1  # 以涨幅为信号，可根据需要修改条件

    # 计算策略收益
    data['Strategy_Return'] = data['Signal'].shift(1) * data['Daily_Return']

    # 计算累计收益
    data['Cumulative_Return'] = (1 + data['Strategy_Return']).cumprod()
    print(data)
    # 绘制累计收益曲线
    plt.figure(figsize=(10, 6))
    plt.plot(data['Cumulative_Return'], label='Strategy Cumulative Return', color='b')
    plt.plot(data['close'] / data['close'].iloc[0], label='Stock Cumulative Return', color='g')
    plt.title("Cumulative Return of Strategy vs. Stock")
    plt.xlabel("Date")
    plt.ylabel("Cumulative Return")
    plt.legend()
    plt.show(block=True)


if __name__ == '__main__':
    # data = getStoneData()
    # # drawKLine(data)
    # drawStroneTrend2(data)
    symbol = "sh601318"
    data = getStoneDataAll(symbol)
    # drawStoneTrend3(symbol, data)
    # backtest(data)
    backtest1(data)
